import torch
from torch.utils.data import Dataset
from torchvision import transforms
import os
from PIL import Image
import numpy as np
import math
import fasttext
class Word_vec_data(Dataset):
    def __init__(self, csv_path,model_path,is_train=True,max_length=50, idx=None):
        line_list = self.get_csv_data(csv_path)
        vec_model = fasttext.load_model(model_path)
        self.data_list=[]
        self.report_id_list=[]
        self.label_list=[]
        for line in line_list:

            if is_train:
                report_id, words, label = line
                word_vec_array = self.get_word_vec_array(vec_model,words,max_length)
                self.data_list.append(word_vec_array)
                self.label_list.append(label)
            else:
                report_id, words = line
                word_vec_array = self.get_word_vec_array(vec_model,words,max_length)
                self.data_list.append((word_vec_array))
            self.report_id_list.append(report_id)
        self.is_train = is_train
        self.report_id_list = np.array(self.report_id_list)
        self.data_list = np.array(self.data_list)
        self.label_list = np.array(self.label_list)
        if isinstance(idx, np.ndarray):
            self.data_list =self.data_list[idx]
            self.report_id_list =self.report_id_list[idx]
            if is_train:
                self.label_list =self.label_list[idx]


    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, index):
        report_id = self.report_id_list[index]
        if self.is_train:
            word_vec_array = self.data_list[index]
            label = self.label_list[index]
            return report_id,word_vec_array,label
        else:
            word_vec_array = self.data_list[index]
            return report_id,word_vec_array
    def get_csv_data(self, csv_path):
        data_list = []

        label_count = {}
        with open(csv_path, 'r') as file:
            for line in file:
                groups = line.strip().split('|,|')
                report_id = groups[0]
                description = groups[1]
                words = description.split(' ')
                if len(groups) >= 3:
                    label = groups[2]
                    areas = label.split(' ')
                    label = np.zeros(17, dtype=np.int)

                    for area_id in areas:
                        if area_id != '':
                            label[int(area_id)] = 1
                            label_count[area_id] = label_count.get(area_id, 0) + 1
                    data_list.append((report_id, words, label))
                else:
                    data_list.append((report_id, words))
        print(sorted(label_count.items(), key=lambda x: x[0]))
        return data_list
    def get_word_vec_array(self,vec_model,words,max_length):
        word_vec_list = []
        for word in words:
            word_vec = vec_model[word]
            word_vec_list.append(word_vec)
        word_vec_tensor = torch.Tensor(word_vec_list)
        n,d = word_vec_tensor.shape
        if n<max_length:
            pad_num = max_length - n
            padding = torch.zeros((pad_num,d))
            word_vec_tensor = torch.cat([word_vec_tensor,padding],dim=0)
        if n>max_length:
            word_vec_tensor = word_vec_tensor[:max_length]

        return word_vec_tensor
if __name__ == "__main__":
    data = TF_IDF_data(r'../data/track1_round1_train_20210222.csv', is_train=True)
    print(data.__getitem__(1))
